package com.newbee.service;

import com.newbee.aop.SkipTrace;
import com.newbee.model.TraceNode;
import com.newbee.model.TraceSearchRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Function;
import java.util.stream.Collectors;

@Service
@Slf4j
public class TraceService {

    @Autowired
    private MongoTemplate mongoTemplate;

    @SkipTrace
    public List<TraceNode> searchTrace(TraceSearchRequest req) {
        Query query = new Query(Criteria.where("traceId").is(req.getTraceId()));
        List<TraceNode> rootNodes = mongoTemplate.find(query, TraceNode.class, "trace_node");
        if (CollectionUtils.isEmpty(rootNodes)) {
            return Collections.emptyList();
        }

        // 1. 构建根节点映射（仅根节点，无需收集子节点，直接用children中的子节点）
        Map<String, TraceNode> rootNodeMap = rootNodes.stream()
                .collect(Collectors.toMap(
                        TraceNode::getNodeId,
                        Function.identity(),
                        (v1, v2) -> v2
                ));

        // 2. 对每个根节点处理：替换末端异步节点（直接用children中的子节点）
        for (TraceNode root : rootNodes) {
            replaceLastAsyncNode(root, rootNodeMap);
        }

        // 3. 筛选并排序根节点
        return rootNodes.stream()
                .filter(Objects::nonNull)
                .filter(t -> !t.isAsync())
                .sorted(Comparator.comparing(TraceNode::getCreateTime))
                .collect(Collectors.toList());
    }

    private static void replaceLastAsyncNode(TraceNode root, final Map<String, TraceNode> traceNodeMap) {
        if (root == null) {
            return;
        }

        // 1. 先递归处理子节点的子分支（确保深层节点先替换）
        CopyOnWriteArrayList<TraceNode> children = root.getChildren();
        if (children == null || children.isEmpty()) {
            return;
        }

        // 2. 遍历子节点（使用迭代器避免并发修改异常）
        for (TraceNode child : children) {
            if (child == null) {
                continue;
            }

            // 3. 先递归处理子节点的下级分支（先深后浅）
            replaceLastAsyncNode(child, traceNodeMap);

            // 4. 判断当前子节点是否为异步节点，需要替换
            if (child.isAsync()) {
                TraceNode replacementNode = traceNodeMap.get(child.getNodeId());
                if (replacementNode != null) {
                    BeanUtils.copyProperties(replacementNode, child);
                    log.info("异步节点[{}]已替换为新节点[{}]", child.getNodeId(), replacementNode.getNodeId());
                } else {
                    log.warn("异步节点[{}]在映射表中未找到对应节点，不替换", child.getNodeId());
                }
            }
        }
    }

}
